import os  # 导入操作系统模块
import matplotlib.pyplot as plt  # 导入绘图模块

import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题


# 定义一个 Plotter 类
class Plotter:
    # 初始化方法，传入参数包括日志目录列表、输出目录、文件名列表、曲线名称列表、坐标轴标签、图例位置和输出图片名
    def __init__(self, log_dirs, output_dir, files, curves, X_label, Y_label, legend, output_png):
        self.log_dirs = log_dirs  # 日志目录列表
        self.output_dir = output_dir  # 输出目录
        # 拼接日志目录和文件名得到完整路径列表
        self.files = [os.path.join(log_dir, file) for log_dir, file in zip(log_dirs, files)]
        self.curves = curves  # 曲线名称列表
        self.X_label = X_label  # X 轴标签
        self.Y_label = Y_label  # Y 轴标签
        self.legend = legend  # 图例位置
        self.output_png = output_png  # 输出图片的文件名
        # 加载所有文件的数据
        self.data = [self.load(file) for file in self.files]

    # 定义加载文件数据的方法
    def load(self, file_path):
        try:
            with open(file_path, 'r') as file:  # 打开文件
                # 读取文件中的每一行，并转换为浮点数，存储在列表中
                num = [float(line.strip()) for line in file.readlines()]
            return num  # 返回数据列表
        except FileNotFoundError:  # 如果文件未找到
            return []  # 返回空列表

    # 定义绘图方法
    def plot(self):
        plt.figure()  # 创建一个新图形

        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False

        # 遍历数据列表和曲线名称列表
        for i, (data, curve) in enumerate(zip(self.data, self.curves)):
            iters = range(len(data))  # 数据的迭代次数
            # 绘制曲线，根据索引 i 设置不同颜色和样式
            if i == 0:
                plt.plot(iters, data, 'red', linewidth=2, label=curve)
            elif i == 1:
                plt.plot(iters, data, 'blue', linewidth=2, label=curve)
            elif i == 2:
                plt.plot(iters, data, 'yellow', linewidth=2, label=curve)
            elif i == 3:
                plt.plot(iters, data, 'green', linewidth=2, label=curve)
            # 可以根据需要添加更多的颜色和样式

        plt.grid(True)  # 显示网格
        plt.xlabel(self.X_label)  # 设置 X 轴标签
        plt.ylabel(self.Y_label)  # 设置 Y 轴标签
        plt.legend(loc=self.legend)  # 设置图例位置

        # 如果输出目录不存在，则创建
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

        # 保存图形到指定的文件
        plt.savefig(os.path.join(self.output_dir, self.output_png), dpi=300, bbox_inches='tight')

        plt.cla()  # 清除当前活动的图形
        plt.close("all")  # 关闭所有图形窗口

# 如果此脚本是主程序，则执行以下代码
if __name__ == "__main__":
    # 设置日志目录列表和输出目录
    log_dirs = [
        r"E:\postgraduate\大论文\seg_train\短数据集 500轮次",
        r"E:\postgraduate\大论文\seg_train\短数据集 500轮次",
        r"E:\postgraduate\大论文\seg_train\长数据集 500轮",
        r"E:\postgraduate\大论文\seg_train\长数据集 500轮"
    ]
    output_directory_C = r"E:\postgraduate\大论文\seg_train\短数据集 500轮次"

    # 设置文件名列表、曲线名称列表、坐标轴标签、图例位置和输出图片名
    files = [
        "f_scores_train.txt",
        "f_scores_val.txt",
        "f_scores_train.txt",
        "f_scores_val.txt"
    ]
    curves = [
        "Ours",
        "UNet+裁剪",
        "UNet+裁剪+深度卷积",
        "UNet"
    ]
    X_label = "epoch"
    Y_label = "f_score"
    legend = "lower right"
    output_png = "对比图3.png"

    # 创建 Plotter 对象，并调用 plot 方法生成图形
    plotter = Plotter(log_dirs, output_directory_C, files, curves, X_label, Y_label, legend, output_png)
    plotter.plot()
